"""
LangGraph 搜索代理工作流

这个文件实现了一个基于LangGraph的智能搜索代理，具有以下特点：

工作流结构：
    START      END
      ↓         ↑ 
    agent ──(条件判断)──→ tools 
      ↑                   ↓              
      └───────────────────┘

循环逻辑：
- agent节点：分析用户输入，决定是否需要调用工具
- 条件判断：检查agent输出是否包含tool_calls
- tools节点：执行工具调用（如搜索）
- 工具执行完成后，结果返回给agent继续处理
- 如果agent不再需要工具，则结束工作流

节点说明：
1. agent节点：调用大语言模型(DeepSeek)，处理用户输入并决定是否需要调用工具
   - 首次调用：分析用户问题，决定是否需要搜索
   - 后续调用：处理工具返回的结果，生成最终回复
2. tools节点：执行具体的工具调用（如搜索功能）
   - 接收agent的工具调用请求
   - 执行搜索操作
   - 将结果返回给agent继续处理
3. 条件边：根据agent的输出决定下一步操作
   - 有tool_calls → 继续到tools节点
   - 无tool_calls → 结束工作流

工作流程：
1. 用户输入消息 → agent节点（首次调用）
2. agent分析消息，决定是否需要工具调用
3. 条件判断：
   - 如果需要工具 → tools节点 → agent节点（循环）
   - 如果不需要工具 → END
4. 在循环中：
   - tools节点执行工具调用
   - 结果返回给agent继续处理
   - agent决定是否需要更多工具调用
5. 当agent不再需要工具时，生成最终回复并结束工作流

状态管理：
- 使用MessagesState管理对话历史
- 使用MemorySaver保存对话状态，支持多轮对话
- 通过thread_id区分不同的对话会话

工具功能：
- search工具：模拟网络搜索，根据查询返回天气信息
- 支持旧金山天气查询的特殊处理

这个工作流展示了LangGraph的基本用法，包括节点定义、条件边、状态管理和工具集成。
"""
import os
os.environ["DEEPSEEK_API_KEY"] = "sk-1a1c5c7e74664f97b9ae47a8e93f2620"
from typing import Annotated, Literal, TypedDict

from langchain_core.messages import HumanMessage
from langchain_deepseek import ChatDeepSeek
from langchain_core.tools import tool
from langgraph.checkpoint.memory import MemorySaver
from langgraph.graph import END, START, StateGraph, MessagesState
from langgraph.prebuilt import ToolNode


# 定义代理要使用的工具
@tool
def search(query: str):
    """调用以浏览网络。"""
    # 这是一个占位符，但不要告诉 LLM...
    if "旧金山" in query.lower() or "san francisco" in query.lower():
        return "温度是 60 华氏度，天气阴霾。"
    return "温度是 90 华氏度，天气晴朗。"


# 定义决定是否继续的函数
def should_continue(state: MessagesState) -> Literal["tools", END]:
    """
    条件路由函数：决定工作流的下一步操作
    
    参数：
        state: MessagesState - 当前对话状态，包含所有消息历史
    
    返回：
        "tools" - 如果最后一条消息包含工具调用，继续到tools节点
        END - 如果没有工具调用，结束工作流
    
    逻辑：
    - 检查最后一条消息是否包含tool_calls
    - 如果有工具调用，说明agent需要执行工具操作
    - 如果没有工具调用，说明agent已经给出最终回复，可以结束
    """
    messages = state['messages']
    last_message = messages[-1]
    # 如果 LLM 发起了工具调用，我们将路由到 "tools" 节点
    if last_message.tool_calls:
        return "tools"
    # 否则，我们停止（回复用户）
    return END


# 定义调用模型的函数
def call_model(state: MessagesState):
    """
    Agent节点函数：调用大语言模型处理消息
    
    参数：
        state: MessagesState - 当前对话状态
    
    返回：
        dict - 包含模型响应的字典，格式为{"messages": [response]}
    
    功能：
    - 获取当前对话历史
    - 调用绑定了工具的DeepSeek模型
    - 模型会根据对话内容决定是否需要调用工具
    - 返回模型的响应消息
    """
    messages = state['messages']
    response = model.invoke(messages)
    # 我们返回一个列表，因为这将被添加到现有列表中
    return {"messages": [response]}


if __name__ == "__main__":
    '''
    START -> agent -(condition)-> tools -> agent
                               -> END
    
    '''
    
    # ==================== 第一步：初始化组件 ====================
    # 定义工具和模型
    tools = [search]
    tool_node = ToolNode(tools)
    # 将 LangChain 工具转换为 OPENAI 工具调用格式
    model = ChatDeepSeek(model="deepseek-chat", temperature=0).bind_tools(tools)

    # ==================== 第二步：构建工作流图 ====================
    # 定义一个新图，使用MessagesState管理状态, 将每个节点的更新合并到状态
    workflow = StateGraph(MessagesState)

    # 添加两个核心节点
    workflow.add_node("agent", call_model)  # 智能代理节点
    workflow.add_node("tools", tool_node)   # 工具执行节点

    # ==================== 第三步：定义图的边和路由 ====================
    # 设置入口点：工作流从agent节点开始
    workflow.add_edge(START, "agent")

    # 添加条件边：实现agent和tools之间的循环
    workflow.add_conditional_edges(
        "agent",        # 从agent节点出发
        should_continue, # 使用条件函数决定路由
    )

    # 添加循环边：tools执行完成后回到agent继续处理
    workflow.add_edge("tools", 'agent')

    # ==================== 第四步：配置状态管理 ====================
    # 初始化内存以在图运行之间保持状态
    checkpointer = MemorySaver()

    # ==================== 第五步：编译和执行 ====================
    # 编译工作流为可执行 LangChain 对象，自动启动 .invoke , .stream , .batch 方法
    app = workflow.compile(checkpointer=checkpointer)

    # 执行工作流
    final_state = app.invoke(
        {"messages": [HumanMessage(content="旧金山的天气如何")]},
        config={"configurable": {"thread_id": 42}}
    )
    '''
    [HumanMessage(content='旧金山的天气如何', additional_kwargs={}, response_metadata={}, id='d2da4225-47ad-4422-9a1e-8cb0c342143c'), 
     AIMessage(content='', additional_kwargs={'tool_calls': [{'id': 'call_0_cddc7941-6fa3-43bd-b0f3-7be65c49e3e6', 'function': {'arguments': '{"query":"旧金山 天气"}', 'name': 'search'}, 'type': 'function', 'index': 0}], 'refusal': None}, response_metadata={'token_usage': {'completion_tokens': 20, 'prompt_tokens': 102, 'total_tokens': 122, 'completion_tokens_details': None, 'prompt_tokens_details': {'audio_tokens': None, 'cached_tokens': 64}, 'prompt_cache_hit_tokens': 64, 'prompt_cache_miss_tokens': 38}, 'model_name': 'deepseek-chat', 'system_fingerprint': 'fp_8802369eaa_prod0623_fp8_kvcache', 'id': 'd41b2d5f-0d49-4531-814f-3c9b64ee7d4c', 'service_tier': None, 'finish_reason': 'tool_calls', 'logprobs': None}, id='run--ac74a2c5-3bfc-46ed-becb-ee48dab4e50c-0', tool_calls=[{'name': 'search', 'args': {'query': '旧金山 天气'}, 'id': 'call_0_cddc7941-6fa3-43bd-b0f3-7be65c49e3e6', 'type': 'tool_call'}], usage_metadata={'input_tokens': 102, 'output_tokens': 20, 'total_tokens': 122, 'input_token_details': {'cache_read': 64}, 'output_token_details': {}}), 
     ToolMessage(content='温度是 60 华氏度，天气阴霾。', name='search', id='de49423e-1011-4dae-96f4-d12d1a6829de', tool_call_id='call_0_cddc7941-6fa3-43bd-b0f3-7be65c49e3e6'), 
     AIMessage(content='旧金山目前的天气是阴霾，温度为60华氏度（约15.6摄氏度）。', additional_kwargs={'refusal': None}, response_metadata={'token_usage': {'completion_tokens': 20, 'prompt_tokens': 157, 'total_tokens': 177, 'completion_tokens_details': None, 'prompt_tokens_details': {'audio_tokens': None, 'cached_tokens': 128}, 'prompt_cache_hit_tokens': 128, 'prompt_cache_miss_tokens': 29}, 'model_name': 'deepseek-chat', 'system_fingerprint': 'fp_8802369eaa_prod0623_fp8_kvcache', 'id': '900467c9-2779-40c7-8543-55cace9b96d6', 'service_tier': None, 'finish_reason': 'stop', 'logprobs': None}, id='run--ddd16cdd-a1e7-4bfb-b0c2-82a10d485e66-0', usage_metadata={'input_tokens': 157, 'output_tokens': 20, 'total_tokens': 177, 'input_token_details': {'cache_read': 128}, 'output_token_details': {}})]
    '''
    print(final_state["messages"][-1].content)

    # 传递相同的 thread_id，对话上下文通过保存的状态得以保留
    final_state = app.invoke(
        {"messages": [HumanMessage(content="纽约的天气如何")]},
        config={"configurable": {"thread_id": 42}}
    )
    for i in final_state["messages"]:
        print(i.content)
    